transformer_engine/common/CMakeLists.txt
option() declared after first use

The option(NVTE_WITH_CUSOLVERMP ...) declaration appears ~12 lines later in the file (at the target_link_libraries block), but this if (NVTE_WITH_CUSOLVERMP) guard relies on it. While this works when -DNVTE_WITH_CUSOLVERMP=ON is passed on the command line (the cache entry is set before CMake processes any option()), a developer who adds a second reference to NVTE_WITH_CUSOLVERMP above the option() call — or who processes the file without the cache — will get confusing behaviour.

CMake convention is to declare option() before first use. Moving the option() declaration to just above the first if (NVTE_WITH_CUSOLVERMP) block (line 281) would improve robustness and readability.

build_tools/utils.py, line 95
FileNotFoundError propagates to the build system

When libtransformer_engine.so cannot be located, common_lib_has_symbol raises FileNotFoundError. The call-site in pytorch.py is a bare if common_lib_has_symbol(...) with no try/except, so any missing-library condition aborts the entire PyTorch extension build rather than gracefully treating it as "symbol absent → no cuSOLVERMp".

This scenario is realistic: the common library might be built in a non-standard path, or a developer may be doing an incremental build of only the PyTorch extension. The function should return False (with an optional warning) instead of raising:

    if lib_path is None:
        import warnings
        warnings.warn(
            "Could not find libtransformer_engine.so; cuSOLVERMp auto-detection disabled.",
            stacklevel=2,
        )
        return False

The subprocess.CalledProcessError path (nm failure) has the same problem and should similarly return False or at most warn.

tests/pytorch/distributed/test_newton_schulz.py, line 252
Missing skip guard for builds without cuSOLVERMp

The file correctly skips when fewer than 2 GPUs are present, but there is no skip when TE was built without NVTE_WITH_CUSOLVERMP=1. In that case, newton_schulz raises RuntimeError("newton_schulz requires Transformer Engine to be built with NVTE_WITH_CUSOLVERMP=1"), which will appear in CI as a test failure rather than a skip.

A simple guard at module level (or inside the test) would mirror the GPU-count check:

import transformer_engine_torch as tex
if not hasattr(tex, "newton_schulz"):
    pytest.skip(
        "Newton-Schulz tests require TE built with NVTE_WITH_CUSOLVERMP=1.",
        allow_module_level=True,
    )

transformer_engine/pytorch/newton_schulz.py, line 813
Undocumented fallback coefficients for num_iterations != 5

When coefficients is None and num_iterations != 5, the function silently falls back to [1.5, -0.5, 0.0] * num_iterations. These are generic cubic Newton-Schulz coefficients and are likely far less accurate than the carefully tuned QUINTIC_COEFFICIENTS. Nothing in the public docstring tells callers about this behaviour.

If 15-iteration runs use the test's custom coefficients (which are different from both defaults), the fallback path is also never exercised by the current tests with the default API — making the generic path untested in practice.

Consider either:

• Documenting that "optimised coefficients are only bundled for num_iterations=5; for other values you must supply coefficients explicitly", and turning the fallback into a ValueError, or
• Bundling a second set of optimised coefficients for the 15-iteration case that matches what the test actually uses.

qa/L1_pytorch_distributed_unittest/test.sh, line 35
Newton-Schulz test runs unconditionally in QA

The test is added without any guard for NVTE_WITH_CUSOLVERMP. If the CI runner is not built with cuSolverMp, this subprocess will fail with a RuntimeError, causing the entire L1 QA job to fail. Other optional features are guarded in their test paths.

Add a conditional:

if [ "${NVTE_WITH_CUSOLVERMP:-0}" == "1" ]; then
    python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py"
fi

tests/pytorch/distributed/test_newton_schulz.py, line 15
Missing skip guard for cuSolverMp availability

The file skips when fewer than 2 GPUs are available, but not when Transformer Engine is built without NVTE_WITH_CUSOLVERMP=1. On such builds, every test fails inside the subprocess with RuntimeError, producing confusing stderr messages.

Add a module-level skip check:

import transformer_engine_torch as tex
if not hasattr(tex, "newton_schulz"):
    pytest.skip(
        "TE not built with NVTE_WITH_CUSOLVERMP=1; skipping Newton-Schulz tests.",
        allow_module_level=True,
    )

build_tools/pytorch.py, line 95
Silent default for CUSOLVERMP_HOME inconsistent with NVSHMEM pattern

The cuSolverMp block silently defaults to "/usr" when CUSOLVERMP_HOME is unset. The NVSHMEM block asserts that NVSHMEM_HOME is explicitly set, providing a clear error message. If the library is not installed under /usr/include and /usr/lib, the build fails with a generic linker error rather than a clear message about the missing environment variable.

Align with the NVSHMEM pattern:

if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))):
    assert (
        os.getenv("CUSOLVERMP_HOME") is not None
    ), "CUSOLVERMP_HOME must be set when compiling with NVTE_WITH_CUSOLVERMP=1"
    cusolvermp_home = Path(os.getenv("CUSOLVERMP_HOME"))
    include_dirs.append(cusolvermp_home / "include")
    library_dirs.append(cusolvermp_home / "lib")
    libraries.append("cusolverMp")
    cxx_flags.append("-DNVTE_WITH_CUSOLVERMP")

transformer_engine/pytorch/newton_schulz.py, line 21
Uses private PyTorch APIs that may change

Lines 20-21 use _get_backend() and _comm_ptr(), which are private PyTorch APIs (underscore prefix indicates internal/unstable). These can change in future PyTorch versions, breaking this code.

Consider using public APIs or documenting this dependency clearly in comments, noting that this code may need updates with new PyTorch releases.

transformer_engine/pytorch/newton_schulz.py, line 67
Use ValueError instead of assert for validation

The assertion on line 65 validates user input. Assertions can be disabled with Python's -O flag, silently allowing invalid inputs. Use ValueError for user-facing validation:

if len(coefficients) != num_iterations * 3:
    raise ValueError(
        f"Unexpected number of coefficients: {len(coefficients)} for {num_iterations} iterations"
    )

transformer_engine/pytorch/newton_schulz.py, line 72
Missing tensor memory layout validation

The C++ code calls data_ptr() which requires contiguous memory. Non-contiguous tensors will cause silent incorrect results. Add a contiguity check before validation:

if not x.is_contiguous():
    raise ValueError("Input tensor must be contiguous (C-order)")

Also add dtype validation since the docstring specifies float32 or bfloat16:

if x.dtype not in (torch.float32, torch.bfloat16):
    raise ValueError(f"Input tensor must be float32 or bfloat16, got {x.dtype}")

tests/pytorch/distributed/test_newton_schulz.py, line 37
Missing subprocess timeout

The distributed test subprocess has no timeout. If it deadlocks (e.g., NCCL communication issue), the test will block indefinitely, hanging the CI job. Add a timeout:

result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False, timeout=300)

transformer_engine/common/CMakeLists.txt, line 317
PUBLIC linkage exposes cuSOLVERMp to all downstream consumers

cuSOLVERMp is linked with PUBLIC visibility, meaning all projects depending on Transformer Engine must have it in their link path, even those that don't use Newton-Schulz. Since the public API doesn't expose cuSOLVERMp types, PRIVATE linkage would provide better encapsulation:

target_link_libraries(transformer_engine PRIVATE ${CUSOLVERMP_LIB})

transformer_engine/common/CMakeLists.txt, line 230
CMake option declared after first use

The option(NVTE_WITH_CUSOLVERMP) is declared at line 308, but used at line 230. Every other optional feature declares the option immediately before using it. While this works when -DNVTE_WITH_CUSOLVERMP=ON is passed on the command line, it breaks the established pattern and could surprise developers adding follow-on logic.

Move the option() declaration to just before line 230.

transformer_engine/common/include/transformer_engine/newton_schulz.h, line 17
Unconditional #include <nccl.h> in public header

newton_schulz.h is installed as a public header. The unconditional #include <nccl.h> forces NCCL into every downstream project's include path, even those with no interest in Newton-Schulz. Since ncclComm_t is only used in functions meaningful when NVTE_WITH_CUSOLVERMP is defined, guarding both the include and declarations would prevent this leakage:

#ifdef NVTE_WITH_CUSOLVERMP
#include <nccl.h>
...
NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank);
void nvte_newton_schulz(...);
#endif

transformer_engine/pytorch/__init__.py, line 62
Unconditional import of optional feature

newton_schulz is unconditionally imported and exported as public API, even when TE is built without NVTE_WITH_CUSOLVERMP. While the function raises a runtime error when called, this exposes the symbol to all users and makes it appear as a supported feature in auto-complete and documentation.

Guard this import similarly to other optional features, or add a note in the docstring that NVTE_WITH_CUSOLVERMP=1 at build time is required.

transformer_engine/common/newton_schulz/newton_schulz.cpp, line 119
Resource leak on exception during context creation

Raw CUDA handles (stream, in_ready, out_ready) are created at lines 99, 102, 104 but not moved into the struct until lines 112-114. If MakeCusolverMpHandle (line 106) or MakeCusolverMpGrid (line 107) throw (both contain NVTE_CHECK_CUSOLVERMP which can throw), these three handles are never cleaned up.

Wrap each handle in an RAII type (similar to CusolverMpHandle and CusolverMpGrid) to ensure safe cleanup on any early-exit path.

build_tools/pytorch.py, line 95
Missing assertion for CUSOLVERMP_HOME

When NVTE_WITH_CUSOLVERMP=1 is set but CUSOLVERMP_HOME is not, this silently falls back to /usr and will produce a confusing linker error at build time rather than a clear configuration message.

Compare to the NVSHMEM_HOME handling directly above (lines 81-83), which explicitly asserts the variable must be set. Add the same guard here:

if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))):
    assert (
        os.getenv("CUSOLVERMP_HOME") is not None
    ), "CUSOLVERMP_HOME must be set when compiling with NVTE_WITH_CUSOLVERMP=1"
    cusolvermp_home = Path(os.getenv("CUSOLVERMP_HOME"))
    ...

qa/L1_pytorch_distributed_unittest/test.sh, line 35
Test runs unconditionally regardless of build config

This test is added unconditionally to the QA script, so it will always execute even when TE is built without NVTE_WITH_CUSOLVERMP=1. The subprocess will fail with a runtime error about missing cuSolverMp support, breaking the CI job.

Add a build-flag guard matching the build configuration:

if [ "${NVTE_WITH_CUSOLVERMP:-0}" = "1" ]; then
    python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py"
fi

tests/pytorch/distributed/test_newton_schulz.py, line 15
No guard for missing cuSolverMp build support

The test only skips when fewer than 2 GPUs are available, but does not check whether TE was built with NVTE_WITH_CUSOLVERMP=1. On a system with 2+ GPUs but a TE build without cuSolverMp, the torchrun subprocess will fail with a runtime error, causing the test to report AssertionError with confusing output rather than a clean skip.

Add an early skip guard:

import transformer_engine_torch as tex

if not hasattr(tex, "newton_schulz"):
    pytest.skip("Newton-Schulz tests require TE built with NVTE_WITH_CUSOLVERMP=1.", allow_module_level=True)

transformer_engine/pytorch/newton_schulz.py, line 64
Fallback coefficients for num_iterations != 5 are undocumented and degrade polynomial degree

When num_iterations != 5 and no custom coefficients are supplied, the fallback is [1.5, -0.5, 0.0] * num_iterations. The trailing 0.0 silently degenerates the quintic polynomial to a cubic one (a·X + b·X³ + 0·X⁵). This means users calling with, e.g., num_iterations=10 will unknowingly use different convergence behavior than the optimized 5-iteration case.

Consider either:

1. Raising a ValueError when num_iterations != 5 and no coefficients are provided, forcing users to supply their own, or
2. Documenting clearly in the docstring that only 5-iteration defaults are optimised and all other counts fall back to generic cubic steps